# -*- coding:utf-8 -*-
import os
import torch
import utils
import numpy as np
import datetime
import time
import logging
import sys
from model.hidden import Hidden
from noise_layers.noiser import Noiser
import pprint
from collections import defaultdict
from average_meter import AverageMeter

def validate(device, this_run_folder, test_run_folder, data_dir, tb_logger=None):
    options_file = os.path.join(this_run_folder, 'options-and-config.pickle')
    train_options, hidden_config, noise_config = utils.load_options(options_file)
    checkpoint, loaded_checkpoint_file_name = utils.load_last_checkpoint(os.path.join(this_run_folder, 'checkpoints'))
    train_options.start_epoch = checkpoint['epoch'] + 1
    if data_dir is not None:
        train_options.train_folder = os.path.join(data_dir, 'train')
        train_options.validation_folder = os.path.join(data_dir, 'val')

    logging.basicConfig(level=logging.INFO,
                        format='%(message)s',
                        handlers=[
                            logging.FileHandler(os.path.join(test_run_folder, f'{train_options.experiment_name}.log')),
                            logging.StreamHandler(sys.stdout)
                        ])
    
    model = Hidden(hidden_config, device, tb_logger)

    assert checkpoint is not None
    logging.info(f'Loading checkpoint from file {loaded_checkpoint_file_name}')
    utils.model_from_checkpoint(model, checkpoint)

    # 将参数信息保存到字符串中
    param_info = ""
    for name, param in model.encoder_decoder.named_parameters():
        param_info += f"Parameter name: {name}, Shape: {param.shape}\n"
        param_info += f"Value: {param}\n"

    # 将参数信息写入到txt文件中
    with open('test_model_params.txt', 'w') as f:
        f.write(param_info)

    print("Model parameters info saved to 'model_params.txt'")
    logging.info('HiDDeN model: {}\n'.format(model.to_stirng()))
    logging.info('Model Configuration:\n')
    logging.info(pprint.pformat(vars(hidden_config)))
    logging.info('\nNoise configuration:\n')
    logging.info(pprint.pformat(str(noise_config)))
    logging.info('\nTraining train_options:\n')
    logging.info(pprint.pformat(vars(train_options)))

    
    # torch.autograd.set_detect_anomaly(True)
    train_data, val_data = utils.get_data_loaders(hidden_config, train_options)

    images_to_save = 8
    saved_images_size = (512, 512)

    epoch_start = time.time()

    first_iteration = True
    validation_losses = defaultdict(AverageMeter)
    logging.info('Running validation for epoch {}/{}'.format(1, train_options.number_of_epochs))
    for image, _ in val_data:
        image = image.to(device)
        message = torch.Tensor(np.random.choice([0, 1], (image.shape[0], hidden_config.message_length))).to(device)
        losses, (encoded_images, noised_images, decoded_messages) = model.validate_on_batch([image, message])
        for name, loss in losses.items():
            validation_losses[name].update(loss)
        if first_iteration:
            if hidden_config.enable_fp16:
                image = image.float()
                encoded_images = encoded_images.float()
            os.makedirs(os.path.join(test_run_folder, 'images'), exist_ok=True)
            revise_image = torch.abs(image - encoded_images) * 5
            revise_image_ = torch.abs(image - noised_images) * 5
            utils.save_images(image.cpu()[:images_to_save, :, :, :],
                                    encoded_images[:images_to_save, :, :, :].cpu(),
                                    noised_images[:images_to_save, :, :, :].cpu(),
                                    revise_image[:images_to_save, :, :, :].cpu(),
                                    revise_image_[:images_to_save, :, :, :].cpu(),
                                    1,
                                    os.path.join(test_run_folder, 'images'), resize_to=saved_images_size)
            first_iteration = False

    utils.log_progress(validation_losses)
    logging.info('-' * 40)
    utils.write_losses(os.path.join(test_run_folder, 'validation.csv'), validation_losses, 1,
                        time.time() - epoch_start)


# def save_images(cover_images, encoder_images, noise_image, epoch, folder, normalize):
#     images = cover_images[:cover_images.shape[0], :, :, :]
#     encoder_images = encoder_images[:encoder_images.shape[0], :, :, :]
#     noise_image = noise_image[:noise_image.shape[0], :, :, :]

#     # scale values to range [0, 1] from original range of [-1, 1]
#     if normalize:
#         images = (images + 1) / 2
#         encoder_images = (encoder_images + 1) / 2
#         noise_image = (noise_image + 1) / 2

#     revise_image = torch.abs(images - encoder_images) * 5
#     revise_image_5 = torch.abs(images - noise_image) * 5

#     stacked_images = torch.cat([images, encoder_images, noise_image, revise_image, revise_image_5], dim=0)
#     filename = os.path.join(folder, 'epoch-{}.png'.format(epoch + 1))
#     # torchvision.utils.save_image(stacked_images, filename=filename, original_images.shape[0], normalize=False)
#     torchvision.utils.save_image(stacked_images, filename, normalize=False)

def main():

    device = torch.device('cuda:2') if torch.cuda.is_available() else torch.device('cpu')
    file_dir = "runs-works-test"
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)

    # 获取当前日期和时间
    current_time = datetime.datetime.now()
    # 格式化日期和时间
    time_str = current_time.strftime("%Y-%m-%d_%H-%M-%S")

    # name = "run" + "_" + "base" + "_" + time_str
    # name = "run" + "_" + "Gaussian-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Quantization-no-normalize" + "_" + time_str
    name = "run" + "_" + "Jpeg-Mask-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Brightness(0.5-0.6)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Contrast(0.5-0.6)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Saturation(0.5-0.6)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Hue(0.2-0.3)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Brightness(1.5-1.6)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Contrast(1.5-1.6)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Saturation(1.5-1.6)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "Hue(（-0.3）-（-0.2）)-no-normalize" + "_" + time_str
    # name = "run" + "_" + "salt-no-normalize" + "_" + time_str

    test_floder = os.path.join(file_dir, name)
    os.makedirs(test_floder, exist_ok=True)

    data_dir = "/home/Share/HiDDeN_Data"     # 测试数据
    
    # 添加要测试的模型
    run_floder = '/home/zzc/Experience/HiDDeN_zzc/runs/run_Jpeg-50-no-normalize_2024-03-08_08-32-37-100'

    validate(device, run_floder, test_floder, data_dir, tb_logger=None)

if __name__ == '__main__':
    main()